#pragma once
#include <torch/torch.h>
#include "TileBounds.h"

using namespace torch::autograd;

//#if defined(USE_HIP) || defined(USE_CUDA) || defined(USE_MPS)

std::tuple<torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor> binAndSortGaussians(int vNumPoints, int vNumIntersects,
        torch::Tensor vXYs,
        torch::Tensor vDepths,
        torch::Tensor vRadii,
        torch::Tensor vCumTilesHit,
        TileBounds vTileBounds);


class CRasterizeGaussians : public Function<CRasterizeGaussians> {
public:
    static torch::Tensor forward(AutogradContext* vCTx,
        torch::Tensor vXYs,
        torch::Tensor vDepths,
        torch::Tensor vRadii,
        torch::Tensor vConics,
        torch::Tensor vNumTilesHit,
        torch::Tensor vColors,
        torch::Tensor vOpacity,
        int vImgHeight,
        int vImgWidth,
        torch::Tensor vBackground);
    static tensor_list backward(AutogradContext* vCTx, tensor_list vGradOutputs);
};

//#endif

